"""数学计算 Agent"""
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage
from typing import TypedDict, Annotated
import operator


class State(TypedDict):
    messages: Annotated[list, operator.add]
    next: str


def create_math_agent(model: ChatOpenAI):
    """创建数学计算 agent"""

    system_prompt = """你是一个专业的数学助手。
你的职责:
1. 解答数学问题(算术、代数、几何等)
2. 进行数学计算和推理
3. 解释数学概念
4. 提供详细的计算步骤

请用清晰、准确的方式回答数学问题。"""

    def math_agent_node(state: State):
        messages = state["messages"]

        # 构建完整的消息历史
        agent_messages = [SystemMessage(content=system_prompt)] + messages

        # 调用模型
        response = model.invoke(agent_messages)

        return {
            "messages": [response],
            "next": "supervisor"
        }

    return math_agent_node